""" Handles IO for files related to the Haardt and Madau 2012 UV Background
http://adsabs.harvard.edu/abs/2012ApJ...746..125H """

import numpy as np
import os.path
import scipy.interpolate
import scipy.integrate

from idle_hands.constants import physical_constants as pc
from idle_hands.constants import hydrogen 
from idle_hands.constants import helium

CC = pc.PhysicalConstants()
H = hydrogen.Hydrogen()
He = helium.Helium()


class Photorates_Table:

    """ Reads HM12 photoionization and heating table """ 

    def __init__( self ):

        # read data
        #-----------------------------------------------------
        local = os.path.dirname(os.path.realpath(__file__))
        fname = local + '/photorates.out'        
        dat = np.loadtxt( fname )
        self._dat = dat

        # organize data
        #-----------------------------------------------------
        self.z = dat[:,0]        
        self.GH1 = dat[:,1]     # H1 photoionization
        self.GH1h = dat[:,2]     # H1 photoheating
        self.GHe1 = dat[:,3]    # He1 photoionization
        self.GHe1h = dat[:,4]    # He1 photoheating
        self.GHe2 = dat[:,5]    # He2 photoionization
        self.GHe2h = dat[:,6]    # He2 photoheating
        self.comptonh = dat[:,7] # compton heating

        self.log1pz = np.log10( 1.0 + self.z )
        
        self.log_GH1 = np.log10( self.GH1 )
        self.log_GH1h = np.log10( self.GH1h )
        self.log_GHe1 = np.log10( self.GHe1 )
        self.log_GHe1h = np.log10( self.GHe1h )
        self.log_GHe2 = np.log10( self.GHe2 )
        self.log_GHe2h = np.log10( self.GHe2h )
        self.log_comptonh = np.log10( self.comptonh )

    def return_values_at_z( self, z ):
        
        log1pz = np.log10( 1.0 + z )

        # find redshift indices
        #-----------------------------------------------------------------
        i_zlo = np.argmin( np.abs( self.log1pz - log1pz ) )
        if self.log1pz[i_zlo] < log1pz:
            i_zhi = i_zlo + 1
        else:
            i_zhi = i_zlo
            i_zlo = i_zlo - 1
            
        dz = (self.log1pz[i_zhi] - self.log1pz[i_zlo])
        z_frac = (log1pz - self.log1pz[i_zlo]) / dz

        dlogGH1 = self.log_GH1[i_zhi] - self.log_GH1[i_zlo]
        logGH1 = self.log_GH1[i_zlo] + dlogGH1 * z_frac
        
        dlogGH1h = self.log_GH1h[i_zhi] - self.log_GH1h[i_zlo]
        logGH1h = self.log_GH1h[i_zlo] + dlogGH1h * z_frac

        dlogGHe1 = self.log_GHe1[i_zhi] - self.log_GHe1[i_zlo]
        logGHe1 = self.log_GHe1[i_zlo] + dlogGHe1 * z_frac
        
        dlogGHe1h = self.log_GHe1h[i_zhi] - self.log_GHe1h[i_zlo]
        logGHe1h = self.log_GHe1h[i_zlo] + dlogGHe1h * z_frac

        dlogGHe2 = self.log_GHe2[i_zhi] - self.log_GHe2[i_zlo]
        logGHe2 = self.log_GHe2[i_zlo] + dlogGHe2 * z_frac
        
        dlogGHe2h = self.log_GHe2h[i_zhi] - self.log_GHe2h[i_zlo]
        logGHe2h = self.log_GHe2h[i_zlo] + dlogGHe2h * z_frac

        dlogcomptonh = self.log_comptonh[i_zhi] - self.log_comptonh[i_zlo]
        logcomptonh = self.log_comptonh[i_zlo] + dlogcomptonh * z_frac

        return ( logGH1, logGH1h, logGHe1, logGHe1h, 
                logGHe2, logGHe2h, logcomptonh )
    


class UVB_Table:

    """ Reads HM12 UVB table """ 

    def __init__( self ):

        # read data
        #-----------------------------------------------------
        local = os.path.dirname(os.path.realpath(__file__))
        fname = local + '/UVB.out'        
        dat = np.loadtxt( fname )

        # organize data
        #-----------------------------------------------------
        self.z = dat[0,:]        

        dat = np.loadtxt( fname, skiprows=21 )
        self._dat = dat
        self.Jnu = dat[:,1:]

        isfinite = self.Jnu > 0
        Jmin = self.Jnu[isfinite].min()
        
        self.logJnu = self.Jnu.copy() 
        self.logJnu[isfinite] = np.log10( self.Jnu[isfinite] )
        self.logJnu[~isfinite] = np.log10( Jmin )

        self.lam_A = dat[:,0]
        self.lam_nm = self.lam_A * 1.0e-1
        self.nu_Hz = CC.c_A_s / self.lam_A
        self.enrg_eV = self.nu_Hz * CC.h_eV_s
        self.enrg_Ry = self.enrg_eV / H.Ry_eV


        self.lam_logA = np.log10( self.lam_A )
        self.lam_lognm = np.log10( self.lam_nm )
        self.nu_logHz = np.log10( self.nu_Hz )
        self.enrg_logeV = np.log10( self.enrg_eV )
        self.enrg_logRy = np.log10( self.enrg_Ry )




class UVB_Standard:

    """ Returns a standard UVB object from HM12 tables. """

    def __init__( self, z, q_min, q_max, Npts=3000, 
                  photoion_fit='verner', verbose=False ): 

        # set input
        #-----------------------------------------------------------------
        self.z = z
        self.q_min = q_min
        self.q_max = q_max

        self.logq_min = np.log10( q_min )
        self.logq_max = np.log10( q_max )

        # attach useful constants
        #-----------------------------------------------------
        self.C = CC
        self.H = H
        self.He = He

        # read HM12 tables
        #-----------------------------------------------------------------
        tab = UVB_Table()

        # convert q min/max (Rydbergs) to Angstroms and Hz
        #-----------------------------------------------------------------
        self.A_min = CC.eV_to_A( q_max * H.Ry_eV )
        self.A_max = CC.eV_to_A( q_min * H.Ry_eV )

        self.logA_min = np.log10( self.A_min )
        self.logA_max = np.log10( self.A_max )

        self.Hz_min = CC.c_A_s / self.A_max
        self.Hz_max = CC.c_A_s / self.A_min

        self.logHz_min = np.log10( self.Hz_min )
        self.logHz_max = np.log10( self.Hz_max )

        # find redshift indices
        #-----------------------------------------------------------------
        i_zlo = np.argmin( np.abs( tab.z - self.z ) )
        if tab.z[i_zlo] < self.z:
            i_zhi = i_zlo + 1
        else:
            i_zhi = i_zlo
            i_zlo = i_zlo - 1
            
        z_frac = (self.z - tab.z[i_zlo]) / (tab.z[i_zhi] - tab.z[i_zlo])


        # make selections on the energy and flux arrays
        # be aware that the raw data is tabulated in Angstroms
        #-----------------------------------------------------------------
        ilo = np.argmin( np.abs( tab.lam_A - self.A_min ) )
        if tab.lam_A[ilo] > self.A_min:
            ilo = ilo-1

        ihi = np.argmin( np.abs( tab.lam_A - self.A_max ) ) + 1
        if tab.lam_A[ihi] < self.A_max:
            ihi = ihi+1
                
        lam_A_clip = tab.lam_A[ilo:ihi]
        lam_logA_clip = np.log10( lam_A_clip )

        dJnu = tab.Jnu[ilo:ihi, i_zhi] - \
               tab.Jnu[ilo:ihi, i_zlo]   
        
        Jnu_clip = tab.Jnu[ilo:ihi, i_zlo] + z_frac * dJnu    

        self.lam_A_clip = lam_A_clip
        self.lam_logA_clip = lam_logA_clip
        self.Jnu_clip = Jnu_clip
        self.logJnu_clip = np.log10( Jnu_clip )

        
        # construct interpolated version between qmin and qmax 
        #-----------------------------------------------------------------    
        self.lam_logA = self.logA_min + \
            np.linspace(0.,1.,Npts) * (self.logA_max-self.logA_min)

        intrpf = scipy.interpolate.interp1d( self.lam_logA_clip, 
                                             self.logJnu_clip )
        self.logJnu = intrpf( self.lam_logA )

        self.lam_A = 10**self.lam_logA
        self.Jnu = 10**self.logJnu


        # now reorder entries so they go from low to high energy
        #-----------------------------------------------------------------    
        self.Jnu_clip = self.Jnu_clip[::-1].copy()
        self.logJnu_clip = self.logJnu_clip[::-1].copy()

        self.Jnu = self.Jnu[::-1].copy()
        self.logJnu = self.logJnu[::-1].copy()
        
        self.lam_A_clip = self.lam_A_clip[::-1].copy()
        self.lam_logA_clip = self.lam_logA_clip[::-1].copy()

        self.lam_A = self.lam_A[::-1].copy()
        self.lam_logA = self.lam_logA[::-1].copy()

        self.nu_Hz_clip = CC.c_A_s / self.lam_A_clip 
        self.enrg_eV_clip = self.nu_Hz_clip * CC.h_eV_s
        self.enrg_Ry_clip = self.enrg_eV_clip / H.Ry_eV


        # now add energy based units
        #-----------------------------------------------------------------    
        self.nu_Hz = CC.c_A_s / self.lam_A 
        self.enrg_eV = self.nu_Hz * CC.h_eV_s
        self.enrg_Ry = self.enrg_eV / H.Ry_eV

        self.nu_logHz = np.log10( self.nu_Hz )
        self.enrg_logeV = np.log10( self.enrg_eV )
        self.enrg_logRy = np.log10( self.enrg_Ry )


        # set photoionization x-sections and integrands
        #-----------------------------------------------------------------
        self.sigmaH1_th = H.photoion_h1( 1.0, fit=photoion_fit )
        self.sigmaH1 = H.photoion_h1( self.enrg_Ry, fit=photoion_fit )

        self.sigmaH1_ratio = self.sigmaH1 / self.sigmaH1_th
        
        # [ps]  = photons / s / cm^2
        # [pcc] = photons / cm^3
        # [GH1] = photons / s
        # [GH1h] = eV / s

        self.dps_dRy  = (4 * np.pi) * self.Jnu / (CC.h_erg_s * self.enrg_Ry)
        self.dpcc_dRy = self.dps_dRy / CC.c
        self.dGH1_dRy = self.dps_dRy * self.sigmaH1
        self.dGH1h_dRy = self.dGH1_dRy * (self.enrg_Ry-1.0)

        # scipy integrate functions take (y,x) NOT (x,y) 
        #-----------------------------------------------------------------
        self.ps_thin  = scipy.integrate.simps( self.dps_dRy,  
                                               self.enrg_Ry ) # photons/s/cm^2

        self.pcc_thin = scipy.integrate.simps( self.dpcc_dRy, 
                                               self.enrg_Ry ) # photons/cm^3

        self.GH1_thin = scipy.integrate.simps( self.dGH1_dRy, 
                                               self.enrg_Ry ) # photoion. rate

        self.GH1h_thin = scipy.integrate.simps( self.dGH1h_dRy,
                                                self.enrg_Ry ) # photoheating

        # calculate grey x-section
        #-----------------------------------------------------------------
        self.grey_sigmaH1 = self.GH1_thin / self.ps_thin

        
        # find energy of photons for grey
        #-----------------------------------------------------------------
        log_sig = np.log10( self.sigmaH1 )
        log_gry = np.log10( self.grey_sigmaH1 )

        i_lo = np.argmin( np.abs( log_sig - log_gry ) )
        if log_sig[i_lo] < log_gry:
            i_hi = i_lo + 1
        else:
            i_hi = i_lo
            i_lo = i_lo - 1
    
        frac = (log_gry - log_sig[i_lo]) / (log_sig[i_hi] - log_sig[i_lo])
        dlogRy = self.enrg_logRy[i_hi] - self.enrg_logRy[i_lo]
        self.grey_enrg_logRy = self.enrg_logRy[i_lo] + frac * dlogRy
        self.grey_enrg_Ry = 10.0**(self.grey_enrg_logRy)




    def return_GH1_thin(self):
        """ Integrates Jnu to calculate the optically thin
        photoionization rate of this spectrum. """

        if self.enrg_Ry.size == 1:
            GH1 = 4 * np.pi * self.Jnu[0] * self.sigmaH1[0]
        else:
            GH1 = scipy.integrate.simps( self.dGH1_dRy, self.enrg_Ry )

        return GH1



    def return_GH1_shld(self, tauH1_th):
        """ Integrates Jnu to calculate the photoionization rate
        after passing through a column with tau(nu=nu_th) = tauHI_th.
        """
        assert isinstance(tauH1_th, float)

        if self.enrg_Ry.size == 1:
            atten = np.exp( -tauH1_th * self.sigmaH1_ratio[0] )
            GH1 = 4 * np.pi * self.Jnu[0] * self.sigmaH1[0]
            GH1 = GH1 * atten
        else:
            atten = np.exp( -tauH1_th * self.sigmaH1_ratio )
            assert atten.size == self.enrg_Ry.size
            GH1 = scipy.integrate.simps( self.dGH1_dRy * atten, self.enrg_Ry )

        return GH1


    def return_GH1h_thin(self):
        """ Integrates Jnu to calculate the optically thin
        photoheating rate of this spectrum (per ion Ry/s). """

        if self.enrg_Ry.size == 1:
            GH1h = 4 * np.pi * self.Jnu[0] * self.sigmaH1[0] * \
                   (self.enrg_Ry[0] - 1.0)
        else:
            GH1h = scipy.integrate.simps( self.dGH1_dRy * (self.enrg_Ry-1.0),
                                         self.enrg_Ry )

        return GH1h


    def return_GH1h_shld(self, tauH1_th):
        """ Integrates Jnu to calculate the photoheating rate
        after passing through a column with tau(nu=nu_th) = tauH1_th.
        (per ion Ry/s)
        """

        assert isinstance(tauH1_th, float)

        if self.enrg_Ry.size == 1:
            atten = np.exp( -tauH1_th * self.sigmaH1_ratio[0] )
            GH1h = 4 * np.pi * self.Jnu[0] * self.sigmaH1[0]
            GH1h = GH1h * atten * (self.enrg_Ry[0] - 1.0)
        else:
            atten = np.exp( -tauH1_th * self.sigmaH1_ratio )
            energy = (self.enrg_Ry-1.0)
            assert atten.size == self.enrg_Ry.size
            GH1h = scipy.integrate.simps( self.dGH1_dRy * atten * energy ,
                                          self.enrg_Ry )

        return GH1h




    def return_pcc_thin(self):
        """ Integrates Jnu to calculate the optically thin
        photons / cm^3 of this spectrum. """

        if self.enrg_Ry.size == 1:
            pcc = 4 * np.pi * self.Jnu[0] / LIGHT
        else:
            pcc = scipy.integrate.simps( self.dpcc_dRy, self.enrg_Ry )

        return pcc
